package org.ricks.common.db;

import org.ricks.common.Pair;
import org.ricks.db.DaoRuntimeException;
import org.ricks.log.Logger;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Collection;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;

/**
 * @author chenwei
 * @Description: 2022年9月9日 13:40:25 发现ORM存在BUG
 * @date 2020/12/1717:39
 */
public class PooledMgr {

    private static DataSource masterDataSource;
    private final ThreadLocal<ConnInfo> masterLocal = ThreadLocal.withInitial(ConnInfo::new);

    private static DataSource slaveDataSource;
    private final ThreadLocal<ConnInfo> slaveLocal = ThreadLocal.withInitial(ConnInfo::new);

    private static class SingleCase {
        public static final PooledMgr INSTANCE = new PooledMgr();
    }

    /**
     * 扫描APP目录找到DataSourceManage 实现类，开始init
     * @param clazz
     * @throws Exception
     */
    public static void init(Class<? extends DaoSourceManage> clazz)  {
        try {
            DaoSourceManage daoSourceManage = clazz.getConstructor().newInstance();
            Pair<DataSource,DataSource> pair = daoSourceManage.initDataSource();
            PooledMgr.masterDataSource = pair.getKey();
            PooledMgr.slaveDataSource = pair.getValue();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static PooledMgr get() {
        return SingleCase.INSTANCE;
    }

    public DataSource getMasterDataSource() {
        return masterDataSource;
    }

    public DataSource getSlaveDataSource() {
        return slaveDataSource;
    }

    public boolean isSlave() {
        return !Objects.isNull(slaveDataSource);
    }

    public PooledMgr() {
        if(masterDataSource == null) {
            Logger.error("检查有没有初始化DataSource！！！");
            throw new DaoRuntimeException("DataSource init failure, %s", getClass().getName());
        }
        if(slaveDataSource == null) {
            Logger.warn("app 未启用db从库,没有初始化从库DataSource ！");
        }
    }

    public Connection getConnection(boolean isMaster)  throws SQLException{
        if (isMaster || Objects.isNull(slaveDataSource)) {
            return getConnection(masterLocal,masterDataSource);
        } else {
            return getConnection(slaveLocal,slaveDataSource);
        }
    }

    public Connection getConnection ( ThreadLocal<PooledMgr.ConnInfo> local,DataSource dataSource ) throws SQLException{
        Connection conn = local.get().getConn();
        if (conn == null) {
            conn = dataSource.getConnection();
            if (local.get().isTransaction()) {
                conn.setAutoCommit(false);
            }
            local.get().setConn(conn);
        }
        return conn;
    }

    private void discardConnection(ThreadLocal<PooledMgr.ConnInfo> local) throws SQLException {
        Connection conn = local.get().getConn();
        if (conn != null) {
            conn.setAutoCommit(true);
            conn.close();
        }
        for (Connection otherConn : local.get().listOtherConn()) {
            otherConn.setAutoCommit(true);
            otherConn.close();
        }
        local.get().dispose();
    }

    public void discardConnectionFromDao(boolean isMaster) throws SQLException {
        if(isMaster || !isSlave()) {
            if (!masterLocal.get().isRemoteDiscard() && !masterLocal.get().isServiceDiscard()) {
                discardConnection(masterLocal);
            }
        } else {
            if (!slaveLocal.get().isRemoteDiscard() && !slaveLocal.get().isServiceDiscard()) {
                discardConnection(slaveLocal);
            }
        }
    }

    private static class ConnInfo {
        private boolean serviceDiscard;
        private boolean remoteDiscard;
        private Connection conn;
        private ConcurrentMap<String, Connection> otherConnMap = new ConcurrentHashMap<>();
        private boolean transaction;
        private ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
        public ConnInfo() {
            dispose();
        }
        public boolean isServiceDiscard() {
            return serviceDiscard;
        }
        public void setServiceDiscard(boolean serviceDiscard) {
            this.serviceDiscard = serviceDiscard;
        }
        public boolean isRemoteDiscard() {
            return remoteDiscard;
        }
        public void setRemoteDiscard(boolean remoteDiscard) {
            this.remoteDiscard = remoteDiscard;
        }
        public Connection getConn() {
            return conn;
        }
        public void setConn(Connection conn) {
            this.conn = conn;
        }
        public boolean isTransaction() {
            return transaction;
        }
        public void setTransaction(boolean transaction) {
            this.transaction = transaction;
        }
        public Connection getOtherConn(String topic) {
            return otherConnMap.get(topic);
        }
        public Connection putOtherConn(String topic, Connection conn) {
            return otherConnMap.putIfAbsent(topic, conn);
        }
        public Collection<Connection> listOtherConn() {
            return otherConnMap.values();
        }
        public void dispose() {
            serviceDiscard = false;
            remoteDiscard = false;
            conn = null;
            otherConnMap.clear();
            transaction = false;
        }
    }
}
